from show_experiment_res import img_tool
from my_indicators import my_indicators
from tqdm import tqdm
import copy
import time
import torch
from easySample import easySample
import shelve
from configure import *
from siamese_graphsage_train import siamese_graphsage_trainer
from cfg_config import malSim_args
from MGMN_train import MGMNTrainer
from siamese_img_main import siamese_img_train
from hash_tlsh import mytlsh
from functionSim_train import functionModelWarpper
from sklearn.metrics import roc_auc_score
import sys
sys.path.append("/home/cyw/projects/function_sim_project/TLSH/")
sys.path.append("/home/cyw/projects/function_sim_project/function_sim/")
sys.path.append("/home/cyw/projects/function_sim_project/Siamese_with_img/")
sys.path.append("/home/cyw/projects/function_sim_project/MGMN")
sys.path.append("/home/cyw/projects/function_sim_project/siamese_Graphsage")

# from functionsim_without_heterogeneous_train import function_without_heterogeous_train_Warpper
# from functionSim_with_cross_interaction_train import function_with_cross_interaction_train_Warpper


modelNames = ["tlsh", "functionSim", "img", "MGMN", "siamese_graphsage",
              "functionSim_with_hete", "functionSim_zero", "angr", "retdec", "radare2", "ida", "EECG"]
experIndic = ["auc", "precision", "recall", "mrr",
              "ndcg", "confusionMatrix", "test_time", "accuracy"]
pairBasePath = r"/home/cyw/projects/function_sim_project/all_data/newPair2/"


class experimentRes():
    """
        用来生成对比试验数据的模型
    """

    def __init__(self, modelName, indicators, pairInf, pairName, indicatorsValue):
        """
            输入:
                加载最佳模型
                训练指标
                训练数据集+训练数据集名称
                训练指标参数
        """
        self.modelSample = easySample()
        self.modelName = modelName
        self.indicators = indicators
        self.myIndicators = my_indicators()
        self.indicatorsValue = indicatorsValue
        # 样本对的正确值
        self.values = []
        for i in range(len(pairInf)):
            self.values.append(pairInf[i][2])
        self.pairInf = pairInf
        self.pairName = pairName
        self.model = None
        self.modelTrainer = None
        self.dataloader = None
        self.get_best_model(modelName)

    def get_best_model(self, modelName):
        """
            加载模型的最佳参数------>model
            部分模型需要加载Trainer--->modelTrain

            functionSim 的消融实验的config需要配置以下,不然这个不好直接修改!!!!!
        """
        if modelName == "tlsh":
            self.model = mytlsh()
        elif modelName == "functionSim":
            # self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_model.pth')
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge.pth')
            # self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge_best_4_8.pth')
            self.modelTrainer = functionModelWarpper(self.model)
        elif modelName == "functionSim_with_hete":
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_with_hete_model.pth')
            self.modelTrainer = functionModelWarpper(self.model)
        elif modelName == "functionSim_zero":
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/functionSim_zero_model.pth')
            self.modelTrainer = functionModelWarpper(self.model)
        elif modelName == "img":
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/img_model.pth')
            # self.model=torch.load(r'/home/cyw/projects/function_sim_project/all_data/models/img_model_copy.pth')
            self.modelTrainer = siamese_img_train()
        elif modelName == "MGMN":
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/MGMN_model.pth')
            self.modelTrainer = MGMNTrainer(self.model, malSim_args)
        elif modelName == "siamese_graphsage":
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/siamese_graphsage_model.pth')
            self.modelTrainer = siamese_graphsage_trainer(self.model)
            self.dataloader = self.modelTrainer.get_model_dataloader(
                self.pairInf, self.pairName, False)
        elif modelName in ["angr", "retdec", "radare2", "ida"]:
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/{}_model.pth'.format(modelName))
            self.modelTrainer = functionModelWarpper(self.model)
        elif modelName == "EECG":
            self.model = torch.load(
                r'/home/cyw/projects/function_sim_project/all_data/models/EECG_model.pth')
            self.modelTrainer = functionModelWarpper(self.model)

    def get_res(self):
        """
            获得结果
            中转站
        """
        time1 = time.time()
        if self.indicators == "auc":
            res = self.auc_res()
        elif self.indicators == "mrr":
            res = self.mrr_res(self.pairName, self.indicatorsValue)
        elif self.indicators == "ndcg":
            res = self.ndcg_res(self.indicatorsValue)
        elif self.indicators == "precision":
            res = self.precision_res_base_auc()
            # res=self.precision_res()
        elif self.indicators == "accuracy":
            res = self.accuracy_res_base_auc()
            # res=self.accuracy_res()
        elif self.indicators == "recall":
            res = self.recall_res(self.indicatorsValue)
        elif self.indicators == "confusionMatrix":
            res = self.confusion_matrix_res()
        elif self.indicators == "test_time":
            res = self.calculate_time_res(self.indicatorsValue)
        time2 = time.time()
        with open("/home/cyw/projects/function_sim_project/all_logs/experiment_res_logs.txt", "a") as file:
            file.write("        结果：{}    耗时：{}\n".format(res, time2-time1))
        print("     结果：{}    耗时：{}\n".format(
            round(res, 5), round(time2-time1, 5)))

    def get_model_calculate_time(self, pairInf, pairName, epochs):
        """
            用于获取各个模型的计算结果的耗时
        """
        print("为了测试结果的合理性，请移除模型中的批处理！！")
        if self.modelName == "tlsh":
            x, y = [], []
            for j in range(epochs):
                for i in range(len(pairInf)):
                    x.append(self.modelSample.get_sample(
                        pairInf[i][0], "tlsh"))
                    y.append(self.modelSample.get_sample(
                        pairInf[i][1], "tlsh"))
            a = time.time()
            score = self.model.score(x, y)
        elif self.modelName == "img":
            imgDatas = self.modelTrainer.get_data(pairInf)
            img_batch_data = self.modelTrainer.trans_pair_into_batch(imgDatas)
            a = time.time()
            for i in range(epochs):
                es, temp = self.modelTrainer.get_scores_from_model(
                    img_batch_data, self.model)
        elif self.modelName == "MGMN" or self.modelName == "siamese_graphsage":
            self.dataloader = self.modelTrainer.get_model_dataloader(
                pairInf, pairName, False)
            a = time.time()
            for i in range(epochs):
                res = self.modelTrainer.get_model_score(self.dataloader)
        else:
            self.dataloader = self.modelTrainer.get_model_dataloader(
                pairInf, pairName, False, False)
            a = time.time()
            for i in range(epochs):
                res = self.modelTrainer.get_model_score(self.dataloader)
        b = time.time()
        return b-a

    def get_model_res(self, pairInf, pairName):
        """
            用于获取各个模型的pairInf的输出结果
            确保一下，pairName并不会加载/
        """
        if self.modelName == "tlsh":
            x, y = [], []
            for i in range(len(pairInf)):
                x.append(self.modelSample.get_sample(pairInf[i][0], "tlsh"))
                y.append(self.modelSample.get_sample(pairInf[i][1], "tlsh"))
            res = self.model.score(x, y)
            # res=self.myIndicators.modify_res(score)
        elif self.modelName == "img":
            imgDatas = self.modelTrainer.get_data(pairInf)
            img_batch_data = self.modelTrainer.trans_pair_into_batch(imgDatas)
            res, temp = self.modelTrainer.get_scores_from_model(
                img_batch_data, self.model)
        elif self.modelName == "MGMN" or self.modelName == "siamese_graphsage":
            self.dataloader = self.modelTrainer.get_model_dataloader(
                pairInf, pairName, False)
            res = self.modelTrainer.get_model_score(self.dataloader)
        else:
            self.dataloader = self.modelTrainer.get_model_dataloader(
                pairInf, pairName, False, False)
            res = self.modelTrainer.get_model_score(self.dataloader)
        return res

    def calculate_time_res(self, epochs):
        """
            模型平均计算一个样本的耗时
            modify-->移除数据加载时间
        """
        res = self.get_model_calculate_time(
            self.pairInf, self.pairName, epochs)
        ans = (res)*1.0/(len(self.pairInf)*epochs)
        print("总的训练耗时为{}".format(res))
        return ans

    def auc_res(self):
        """
            计算auc值
        """
        res = self.get_model_res(self.pairInf, self.pairName)
        aucValue = self.myIndicators.auc(self.values, res)
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/auc/{}_auc".format(self.modelName)) as file:
            file["pred"] = res
            file["true"] = self.values
        return aucValue

    # 下面三个函数,实现都比较取巧
    # 都直接使用MRR的计算结果
    def recall_res(self, recall_k):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list = file["ranked_list"]
            ground_truth = file["ground_truth"]
        res = 0.0
        for i in range(len(ranked_list)):
            res += self.myIndicators.recall_at_k(
                ranked_list[i], ground_truth[i], recall_k)
        ans = res/len(ranked_list)
        return ans

    def ndcg_res(self, ndcg_k):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list = file["ranked_list"]
            ground_truth = file["ground_truth"]
        res = 0.0
        for i in range(len(ranked_list)):
            res += self.myIndicators.ndcg_at_k(
                ranked_list[i], ground_truth[i], ndcg_k)
        ans = res/len(ranked_list)
        return ans

    def precision_res(self):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list = file["ranked_list"]
            ground_truth = file["ground_truth"]
        res = 0.0
        for i in range(len(ranked_list)):
            res += self.myIndicators.precision(ranked_list[i], ground_truth[i])
        ans = res/len(ranked_list)
        return ans

    def accuracy_res(self):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list = file["ranked_list"]
            ground_truth = file["ground_truth"]
        res = 0.0
        for i in range(len(ranked_list)):
            res += self.myIndicators.accuracy(ranked_list[i], ground_truth[i])
        ans = res/len(ranked_list)
        return ans

    def is_same_fam(self, x, y):
        """
            用于判断各个模型的分类结果是否属于同一个家族
        """
        limit_num = 0.8
        VALUE_TEMP = 0.5
        if self.modelName == "tlsh":
            if (x >= 0.5 and y > VALUE_TEMP) or (x < 0.5 and y < VALUE_TEMP):
                return True
        elif self.modelName == "img":
            if (x >= limit_num and y > VALUE_TEMP) or (x < limit_num and y < VALUE_TEMP):
                return True
        elif self.modelName == "MGMN":
            if (x >= limit_num and y > VALUE_TEMP) or (x < limit_num and y < VALUE_TEMP):
                return True
        elif self.modelName in ["functionSim", "functionSim_zero", "functionSim_with_hete"]:
            if (x >= limit_num and y > VALUE_TEMP) or (x < limit_num and y < VALUE_TEMP):
                return True
        elif self.modelName == "siamese_graphsage":
            if (x >= limit_num and y > VALUE_TEMP) or (x < limit_num and y < VALUE_TEMP):
                return True
        return False

    def accuracy_res_base_auc(self):
        res = self.get_model_res(self.pairInf, self.pairName)
        size = len(self.values)
        tp, tn, fp, fn = 0, 0, 0, 0
        for i in range(size):
            tar = self.values[i]
            temp = res[i]
            if self.is_same_fam(temp, tar):
                if tar > 0.5:
                    tp += 1
                else:
                    tn += 1
            else:
                if tar < 0.5:
                    fn += 1
                else:
                    fp += 1
        return (tp+tn)*1.0/(tp+fp+tn+fn)

    def precision_res_base_auc(self):
        res = self.get_model_res(self.pairInf, self.pairName)
        size = len(self.values)
        tp, tn, fp, fn = 0, 0, 0, 0
        for i in range(size):
            tar = self.values[i]
            temp = res[i]
            if self.is_same_fam(temp, tar):
                if tar > 0.5:
                    tp += 1
                else:
                    tn += 1
            else:
                if tar < 0.5:
                    fn += 1
                else:
                    fp += 1
        return tp*1.0/(tp+fn)

    def confusion_matrix_res(self):
        with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
            ranked_list = file["ranked_list"]
            ground_truth = file["ground_truth"]
            sim_res = file["sim_res"]
            name_list = file["nameList"]
            name_to_lable = file["name_to_lable"]
            lable_to_name = file["lable_to_name"]

        # 是乱序的，这里按照家族来合并一下
        lth = len(lable_to_name)  # 家族的个数
        lable_to_ind = {}
        name_to_ind = {}
        cnt = 0
        lable_nums = {}
        lables = []
        for i in lable_to_name.keys():
            lables.append(i)
            lable_to_ind[i] = cnt
            cnt += 1
            lable_nums[i] = len(lable_to_name[i])*(len(lable_to_name[i])-1)
        cnt = 0
        for i in range(len(name_list)):
            name_to_ind[name_list[i]] = cnt  # 这份是啥，我忘记了
            cnt += 1

        # 计算混淆矩阵
        # 同时保存误报、漏报信息
        loubao = {}
        matrix_res = [[0] * (lth+1) for i in range(lth+1)]
        for family in lables:
            a = lable_to_ind[family]
            for name in lable_to_name[family]:
                ind = name_to_ind[name]
                x = ranked_list[ind]
                size = len(ground_truth[ind])
                for i in range(size):
                    tar = x[i]
                    b = lable_to_ind[name_to_lable[tar]]
                    matrix_res[a+1][b+1] += 1
                    if a != b:
                        if family not in loubao:
                            loubao[family] = [(name, tar, name_to_lable[tar])]
                        else:
                            loubao[family].append(
                                (name, tar, name_to_lable[tar]))
        for i in range(1, lth+1):
            matrix_res[0][i] = lables[i-1]
            matrix_res[i][0] = lables[i-1]
        print("正在绘制表格")
        gene_img = img_tool()
        gene_img.draw_confusion_matrix(self.modelName, matrix_res, lables)
        print("表格绘制完成")
        with open("/home/cyw/projects/function_sim_project/all_data/indicators/loubao/loubao.txt", "a") as file:
            file.write("{}  的漏报信息：\n".format(self.modelName))
            for i in loubao.keys():
                temp = {}
                for originName, name, fam in loubao[i]:
                    if fam not in temp:
                        temp[fam] = [originName+"--"+name]
                    else:
                        temp[fam].append(originName+"--"+name)
                file.write("    家族名：{}\n".format(i))
                for j in temp.keys():
                    file.write("        误判家族名：{}\n".format(j))
                    res = " ".join(temp[j])
                    file.write("            {}\n".format(res))
        print("数据保存成功")
        return True

    def mrr_res(self, saveName, use_save_data):
        """
            一个样本和所有的样本进行相似性的度量,然后保留对与对之间的相似值
            计算出每个家族所包含的所有样本,做成list,当成ground Truth

            use_save_data,表示是否使用之前保存的数据

        """
        if use_save_data == "use_save_data":
            with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
                ranked_list = file["ranked_list"]
                ground_truth = file["ground_truth"]
        else:
            if gene_disassemble_pair == False:
                # with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/pair_infs/{}".format(saveName)) as file:
                # with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/{}".format(saveName)) as file:
                with shelve.open(pairBasePath+"{}".format(saveName)) as file:
                    allPairs = file["pairs"]
                    lable_to_name = file["lable_to_name"]
                    name_to_lable = file["name_to_lable"]
                    ans = []
            else:
                with shelve.open(r"/home/cyw/projects/function_sim_project/all_data/pair_infs/dis_{}".format(saveName)) as file:
                    allPairs = file["pairs"]
                    lable_to_name = file["lable_to_name"]
                    name_to_lable = file["name_to_lable"]
                    ans = []

            # 每一个pairs就是一个样本和其他所有样本的组合
            for pairs in tqdm(allPairs):
                temp = []
                # 这个应该是test
                res = self.get_model_res(pairs, self.pairName)
                for i in range(len(res)):
                    temp.append((res[i], pairs[i][0], pairs[i][1]))
                # 按照相似值的大小排序
                temp.sort(reverse=True)
                ans.append(copy.deepcopy(temp))

            # 计算mrr值
            ranked_list = []
            ground_truth = []
            nameList = []
            for i in range(len(ans)):
                temp = []
                for j in range(len(ans[i])):
                    temp.append(ans[i][j][2])
                ranked_list.append(temp)
                ground_truth.append(lable_to_name[name_to_lable[ans[i][0][1]]])
                nameList.append(ans[i][0][1])

            #   保存计算数据
            with shelve.open("/home/cyw/projects/function_sim_project/all_data/indicators/MRR/all_to_all_{}".format(self.modelName)) as file:
                file["ranked_list"] = ranked_list
                file["ground_truth"] = ground_truth
                file["nameList"] = nameList
                # res没用。。。这里只保存了最后一行的相似值
                file["sim_res"] = res
                file["lable_to_name"] = lable_to_name
                file["name_to_lable"] = name_to_lable

        return self.myIndicators.MRR(ranked_list, ground_truth)


def getPairInf(dataName, path):
    with shelve.open(path) as file:
        if dataName == "valid":
            testPair = file["validPair"]
        elif dataName == "test":
            testPair = file["testPair"]
        # 移除家族信息
        pairInf = [row[:3] for row in testPair]
    return pairInf


def run_for_res(modelNames, indics):
    """ 
        用来获得各个模型的各个指标值

        输入:[模型名],[指标名],以及对应指标的参数
        输出:无,保存到日志文件中
    """
    dataName = "test"
    # dataName="valid"

    # 获得用于生成指标的数据对
    if gene_disassemble_pair == False:
        # pairInf = getPairInf(dataName,"/home/cyw/projects/function_sim_project/all_data/pair_infs/pair_infs")
        pairInf = getPairInf(dataName, pairBasePath+"pair_infs")
        # pairInf = getPairInf(dataName,"/home/cyw/projects/function_sim_project/all_data/EECG_pair_infs/pair_infs")
    else:
        # 加载反汇编实验所需的样本对
        pairInf = getPairInf(
            dataName, "/home/cyw/projects/function_sim_project/all_data/pair_infs/dis_pair_infs")

    for modelName in modelNames:
        if modelName in ["functionSim_zero", "functionSim_with_hete"]:
            print("functionSim消融实验，注意配置信息是否切换！！！")
        for indic in indics.keys():
            # 部分指标有参数需要设置
            value = indics[indic]
            with open("/home/cyw/projects/function_sim_project/all_logs/experiment_res_logs.txt", "a") as file:
                file.write("{}模型  计算{}值:\n".format(modelName, indic))
            print("{}模型  计算{}值:\n".format(modelName, indic))
            tools = experimentRes(modelName, indic, pairInf, dataName, value)
            tools.get_res()


# 运行命令:nohup python /home/cyw/projects/function_sim_project/basic_script/experiment_res.py
# 消融实验部分代码逻辑未完成,模型指标能放跑通未验证
if __name__ == "__main__":
    """
        评价指标:
            auc         √
            precision   √
            recall@k    √
            mrr         √
            ndcg@k      √
            testTime    √（感觉不太正确，部分模型样本加载太慢了）        
    """

    # modelName=["siamese_graphsage"]
    # modelName=["functionSim_zero"]
    # modelName=["functionSim_with_hete"]
    # modelName=["functionSim"]
    # /home/cyw/projects/function_sim_project/all_data/models/functionSim_model_cross_edge.pth
    modelName = ["MGMN"]
    # modelName=["img"]
    # modelName=["tlsh"]
    # modelName=["EECG"]
    # modelName=["tlsh","img","MGMN","functionSim","siamese_graphsage"]
    # modelName=["functionSim","siamese_graphsage"]
    # modelName = ["angr","retdec","radare2","ida"]
    # modelName = ["angr"]
    # modelName = ["retdec"]
    # modelName = ["ida"]
    # modelName = ["radare2"]
    indic = {
        "auc": True,
        "mrr": "else",
        # "mrr":"use_save_data",
        # "precision":True,
        # "accuracy":True,
        "ndcg": 10,
        # "recall":10,
        # "test_time":10
    }
    is_able_run = True
    for name in modelName:
        if name not in modelNames:
            is_able_run = False
    for ind in indic.keys():
        if ind not in experIndic:
            is_able_run = False
    if is_able_run:
        run_for_res(modelName, indic)
    else:
        print("模型或者指标指定错误！！！")
